"""
实现召回
"""
import json
import os.path

import chatbot.config as config
from sklearn.feature_extraction.text import TfidfVectorizer
import pysparnn.cluster_index as ci
import pickle

from chatbot.dnn.recall.BM25Vectorizer import Bm25Vectorizer
from chatbot.dnn.recall.fasttext_vectors import FastTextVectorizer


class Sentence2Vector:
    def __init__(self, by_word=False, method="bm25"):
        self.qa_dict = json.load(open(config.recall_qa_dict_path, encoding="utf-8"))
        self.by_word = by_word
        self.index_path = config.recall_search_index_path
        if method == "bm25":
            self.index_path = config.recall_search_index_path + "." + method
            self.vectorizer = Bm25Vectorizer()
        elif method == "fasttext":
            self.index_path = config.recall_search_index_path + "." + method
            self.vectorizer = FastTextVectorizer()
        else:
            self.vectorizer = TfidfVectorizer()

    def build_vectors(self):
        lines = [q for q in self.qa_dict]
        if self.by_word:
            lines_cuted = [" ".join(self.qa_dict[q]["q_cut_by_word"]) for q in lines]
        else:
            lines_cuted = [" ".join(self.qa_dict[q]["q_cut"]) for q in lines]

        features_vec = self.vectorizer.fit_transform(lines_cuted)
        search_index = self.get_cp(features_vec, lines)
        return self.vectorizer, features_vec, lines_cuted, search_index

    def get_cp(self, vectors, data):
        if os.path.exists(self.index_path):
            search_index = pickle.load(open(self.index_path, "rb"))
        else:
            search_index = self.build_cp(vectors, data)
        return search_index

    def build_cp(self, vectors, data):
        search_index = ci.MultiClusterIndex(vectors, data)
        pickle.dump(search_index, open(self.index_path, "wb"))
        return search_index
